import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import os
from scipy import stats
import matplotlib as mpl

# 设置全局中文字体 [^优化]
def set_chinese_font():
    plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'sans-serif']
    plt.rcParams['axes.unicode_minus'] = False
    return plt.rcParams

set_chinese_font()

def baseline_comparison(df):
    """组间对比（优化统计方法）[8](@ref)"""
    if 'Group' not in df.columns:
        df['Group'] = 'Unknown'
    
    # 仅处理数值特征
    num_features = df.select_dtypes(include=np.number).columns
    comparison = pd.DataFrame()
    
    for feature in num_features:
        asd_data = df[df['Group']=='ASD'][feature].dropna()
        td_data = df[df['Group']=='TD'][feature].dropna()
        
        if len(asd_data) < 3 or len(td_data) < 3:
            continue
            
        # 非参数检验 [2](@ref)
        u_stat, p_value = stats.mannwhitneyu(asd_data, td_data)
        # 效应量计算 [8](@ref)
        n1, n2 = len(asd_data), len(td_data)
        pooled_std = np.sqrt(((n1-1)*asd_data.std()**2 + (n2-1)*td_data.std()**2) / (n1+n2-2))
        cohens_d = (asd_data.mean() - td_data.mean()) / pooled_std if pooled_std != 0 else 0
        
        comparison.at[feature, 'ASD_Mean'] = asd_data.mean()
        comparison.at[feature, 'ASD_Std'] = asd_data.std()
        comparison.at[feature, 'TD_Mean'] = td_data.mean()
        comparison.at[feature, 'TD_Std'] = td_data.std()
        comparison.at[feature, 'P_Value'] = p_value
        comparison.at[feature, 'Cohens_d'] = cohens_d
    
    return comparison.sort_values(by='P_Value')

def plot_feature_distributions(df, save_path=''):
    """特征分布图（中文标签）[^优化]"""
    if 'Group' not in df.columns:
        df['Group'] = 'Unknown'
    
    num_features = df.select_dtypes(include=np.number).columns
    
    for feature in num_features:
        plt.figure(figsize=(10, 6))
        sns.kdeplot(data=df, x=feature, hue='Group', 
                    fill=True, common_norm=False, alpha=0.5)
        plt.title(f'{feature} 组间分布')
        plt.tight_layout()
        
        if save_path:
            output_file = os.path.join(save_path, f'{feature}_distribution.png')
            plt.savefig(output_file, dpi=300, bbox_inches='tight')
        plt.close()

def plot_correlation_heatmap(df, save_path=''):
    """相关性热力图（中文标签）[^优化]"""
    num_df = df.select_dtypes(include=np.number)
    if num_df.empty: return
        
    corr_matrix = num_df.corr()
    plt.figure(figsize=(12, 10))
    mask = np.triu(np.ones_like(corr_matrix, dtype=bool))
    sns.heatmap(corr_matrix, mask=mask, annot=True, 
               cmap='coolwarm', vmin=-1, vmax=1,
               fmt=".2f", annot_kws={"fontsize":9})
    plt.title('特征相关性矩阵')
    
    if save_path:
        output_file = os.path.join(save_path, 'feature_correlations.png')
        plt.savefig(output_file, dpi=300, bbox_inches='tight')
    plt.close()